import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import pandas as pd
from sklearn.model_selection import train_test_split
from string import printable
import numpy as np

from t1 import MaliciousURLDetector, URLDataset
import random

seed=40
torch.manual_seed(seed) # 为CPU设置随机种子
torch.cuda.manual_seed(seed) # 为当前GPU设置随机种子
torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU，为所有GPU设置随机种子
np.random.seed(seed)  # Numpy module.
random.seed(seed)  # Python random module.	
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True

def evaluate_metrics(model, dataloader, device):
    model.eval()
    true_labels = []
    predicted_labels = []

    with torch.no_grad():
        for inputs, labels in dataloader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            predicted = (outputs > 0.5).float()

            true_labels.extend(labels.cpu().numpy())
           # predicted_labels.extend(predicted.cpu().numpy())
            predicted_labels.extend(predicted.squeeze().cpu().numpy())
            #predicted_labels.extend(predicted.squeeze().cpu().numpy())
            #predicted_value = predicted.item()  # 如果predicted是0维张量，使用.item()将其转换为Python标量  
            #predicted_labels.append(predicted_value)  # 使用append而不是extend

    true_labels = np.array(true_labels)
    predicted_labels = np.array(predicted_labels)
    print(true_labels)
    print(predicted_labels)
    
    TP = np.sum((true_labels == 1) & (predicted_labels == 1))
    TN = np.sum((true_labels == 0) & (predicted_labels == 0))
    FP = np.sum((true_labels == 0) & (predicted_labels == 1))
    FN = np.sum((true_labels == 1) & (predicted_labels == 0))
    '''
    TP = np.sum((true_labels == 0) & (predicted_labels == 0))
    TN = np.sum((true_labels == 1) & (predicted_labels == 1))
    FP = np.sum((true_labels == 1) & (predicted_labels == 0))
    FN = np.sum((true_labels == 0) & (predicted_labels == 1))
    '''
    accuracy = (TP + TN) / (TP + TN + FP + FN)
    precision = TP / (TP + FP) if TP + FP > 0 else 0
    recall = TP / (TP + FN) if TP + FN > 0 else 0
    '''
    accuracy = (TP.astype(float) + TN.astype(float)) / (TP.astype(float) + TN.astype(float) + FP.astype(float) + FN.astype(float))  
    precision = TP.astype(float) / (TP.astype(float) + FP.astype(float)) if (TP.astype(float) + FP.astype(float)) > 0 else 0  
    recall = TP.astype(float) / (TP.astype(float) + FN.astype(float)) if (TP.astype(float) + FN.astype(float)) > 0 else 0
    '''
    f1 = 2 * (precision * recall) / (precision + recall) if precision + recall > 0 else 0

    return accuracy, precision, recall, f1


if __name__ == '__main__':
    # 加载数据集
    df = pd.read_csv('D:\\urldata.csv',encoding="ANSI")
    urls = df['url'].tolist()
    labels = df['isMalicious'].tolist()
    max_len = 75
    vocab_size = len(printable) + 1

    # 数据集划分
    batch_size=32
    train_urls, test_urls, train_labels, test_labels = train_test_split(urls, labels, test_size=0.2, random_state=33)
    train_dataset = URLDataset(train_urls, train_labels, max_len)
    test_dataset = URLDataset(test_urls, test_labels, max_len)
    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_dataloader = DataLoader(test_dataset, batch_size=batch_size)

    # 加载模型
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    emb_dim = 32
    lstm_hidden_dim = 64
    num_classes = 1
    model = MaliciousURLDetector(max_len, vocab_size, emb_dim, lstm_hidden_dim, num_classes).to(device)
    # 获取当前模型的参数  
    current_state_dict = model.state_dict()  
    
    # 加载已保存的模型参数  
    #saved_state_dict = torch.load('F:\\model-pinjiecnn+bilstm+zhuyili+dropout+shuangchihua_epoch_9.pth')
    

    # 导入训练好的模型权重
    #model.load_state_dict(torch.load('F:\\model-zuiyuanshi.pth'))
    model.load_state_dict(torch.load('F:\\zuimodel-5.pth'))
    
    # 评估模型性能
    accuracy, precision, recall, f1 = evaluate_metrics(model, test_dataloader, device)
    print(f'Accuracy: {accuracy:.4f}')
    print(f'Precision: {precision:.4f}')
    print(f'Recall: {recall:.4f}')
    print(f'F1 Score: {f1:.4f}')